import os

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'


class ConfigurableComicGenerator:
    """
    🎨 漫画生成器使用说明：

    限制说明：
    - 按句号分割场景，每个句号生成一张图片
    - 每个场景建议15-30字，描述要具体明确
    - 一次生成3-8个场景为宜，过多可能内存不足
    - 仅支持预设角色：小明、小红、老师、小猫、小狗等
    - 新角色需手动添加到prompt_mapping字典

    最佳实践：
    - ❌ 避免："小明很高兴"
    - ✅ 推荐："小明开心地跳起来，脸上露出灿烂的笑容"
    - ❌ 避免复杂场景："公园里很多人做各种活动"
    - ✅ 推荐单一场景："小明在开满鲜花的公园里散步"
    """

    def __init__(self):
        self.model = None

    def load_model(self):
        """加载AI模型 - 需要6GB+ GPU显存或16GB+系统内存"""
        try:
            import torch
            from diffusers import DiffusionPipeline

            print("🔄 正在加载模型...")

            if torch.cuda.is_available():
                print("🎯 使用GPU加速 - 每图约10-30秒")
                self.model = DiffusionPipeline.from_pretrained(
                    "cagliostrolab/animagine-xl-3.1",
                    torch_dtype=torch.float16,
                    use_safetensors=True,
                    safety_checker=None
                ).to("cuda")
            else:
                print("💻 使用CPU模式 - 每图约2-5分钟")
                self.model = DiffusionPipeline.from_pretrained(
                    "cagliostrolab/animagine-xl-3.1",
                    torch_dtype=torch.float32,
                    safety_checker=None
                )

            self.model.enable_attention_slicing()
            print("✅ 模型加载成功")
            return True

        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            return False

    def generate_from_text(self, novel_text, output_dir="my_comic"):
        """
        从文本生成漫画

        参数：
        novel_text: 故事文本，按句号分割场景
        output_dir: 输出文件夹名称

        注意：每个句号生成一张图片，建议3-8个场景
        """
        if not self.model:
            print("❌ 请先加载模型")
            return False

        try:
            import torch
            from PIL import Image, ImageDraw, ImageFont
            import time

            # 创建输出目录
            os.makedirs(output_dir, exist_ok=True)

            # 🎯 关键：按句号分割场景 - 每个句号生成一张图片
            scenes = []
            sentences = novel_text.strip().split('。')
            for i, sentence in enumerate(sentences):
                if sentence.strip():
                    scenes.append({
                        "text": sentence.strip() + '。',
                        "prompt": self._create_prompt(sentence.strip())
                    })

            print(f"📖 分析出 {len(scenes)} 个场景")

            # ⚠️ 警告：场景过多可能内存不足
            if len(scenes) > 10:
                print("⚠️  场景数量较多，建议控制在10个以内")

            # 生成每个场景
            for i, scene in enumerate(scenes):
                print(f"\n🎨 生成场景 {i + 1}: {scene['text']}")
                print(f"  提示词: {scene['prompt']}")

                start_time = time.time()

                try:
                    with torch.no_grad():
                        image = self.model(
                            scene["prompt"],
                            width=512,  # 固定分辨率
                            height=512,  # 固定分辨率
                            num_inference_steps=25,
                            guidance_scale=7.5,
                            negative_prompt="ugly, blurry, low quality, deformed, distorted, bad anatomy"
                        ).images[0]

                    # 直接生成带文字的漫画面板，不保存中间图片
                    panel_path = f"{output_dir}/panel_{i + 1:02d}.png"
                    success = self._create_comic_panel(image, scene, panel_path)

                    generation_time = time.time() - start_time
                    if success:
                        print(f"✅ 完成 ({generation_time:.1f}s) - {panel_path}")
                    else:
                        print(f"⚠️  生成完成但添加字幕失败 ({generation_time:.1f}s) - {panel_path}")

                except Exception as e:
                    print(f"❌ 场景 {i + 1} 生成失败: {e}")
                    continue

            print(f"\n🎉 漫画生成完成！查看: {output_dir}")
            return True

        except Exception as e:
            print(f"❌ 生成过程出错: {e}")
            return False

    def _create_comic_panel(self, image, scene, output_path):
        """直接创建带文字的漫画面板"""
        try:
            from PIL import Image, ImageDraw, ImageFont

            # 创建漫画面板（更大的画布来容纳图像和文字）
            panel_width, panel_height = 600, 800
            panel = Image.new('RGB', (panel_width, panel_height), '#F0F0F0')

            # 调整图像大小并粘贴到面板中
            img_resized = image.resize((panel_width - 40, 600), Image.LANCZOS)
            panel.paste(img_resized, (20, 20))

            draw = ImageDraw.Draw(panel)

            # 添加漫画边框
            draw.rectangle([4, 4, panel_width - 4, panel_height - 4], outline='#222222', width=4)
            draw.rectangle([8, 8, panel_width - 8, panel_height - 8], outline='#444444', width=2)

            # 尝试加载中文字体
            font = self._load_chinese_font(20)  # 增大字体大小

            # 添加文字区域背景
            text_bg_height = 140  # 增加文字区域高度
            draw.rectangle([10, panel_height - text_bg_height - 10,
                            panel_width - 10, panel_height - 10],
                           fill='white', outline='#CCCCCC', width=2)

            # 添加文字
            text = scene["text"]
            if font:
                # 使用改进的文本换行
                text_lines = self._wrap_chinese_text(text, font, panel_width - 40)

                # 计算文字起始位置
                text_y = panel_height - text_bg_height + 20
                for line in text_lines:
                    bbox = draw.textbbox((0, 0), line, font=font)
                    line_width = bbox[2] - bbox[0]
                    line_x = (panel_width - line_width) // 2
                    # 添加文字阴影效果，提高可读性
                    draw.text((line_x + 1, text_y + 1), line, fill='#666666', font=font)
                    draw.text((line_x, text_y), line, fill='#000000', font=font)
                    text_y += 35
            else:
                # 如果没有字体，至少保存图片
                print("⚠️  无法加载中文字体，图片将不带字幕")

            # 保存面板
            panel.save(output_path, quality=95)
            return True

        except Exception as e:
            print(f"❌ 创建漫画面板失败: {e}")
            # 如果失败，直接保存原图
            try:
                image.save(output_path)
                return True
            except:
                return False

    def _load_chinese_font(self, font_size=20):
        """加载中文字体"""
        from PIL import ImageFont

        font_paths = [
            "C:/Windows/Fonts/simhei.ttf",  # Windows 黑体
            "C:/Windows/Fonts/simsun.ttc",  # Windows 宋体
            "C:/Windows/Fonts/msyh.ttc",  # Windows 微软雅黑
            "C:/Windows/Fonts/msyhbd.ttc",  # Windows 微软雅黑粗体
            "C:/Windows/Fonts/simkai.ttf",  # Windows 楷体
            "/System/Library/Fonts/PingFang.ttc",  # macOS
            "/Library/Fonts/Arial Unicode.ttf",  # macOS
            "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # Linux
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"  # Linux 回退
        ]

        for path in font_paths:
            try:
                if os.path.exists(path):
                    return ImageFont.truetype(path, font_size)
            except:
                continue

        # 如果所有字体都加载失败，尝试使用默认字体
        try:
            return ImageFont.load_default()
        except:
            return None

    def _wrap_chinese_text(self, text, font, max_width):
        """中文文本换行 - 改进版本"""
        lines = []
        current_line = ""

        for char in text:
            test_line = current_line + char
            # 使用更精确的文本宽度计算
            bbox = font.getbbox(test_line) if hasattr(font, 'getbbox') else (
            0, 0, len(test_line) * font.size, font.size)
            text_width = bbox[2] - bbox[0] if hasattr(font, 'getbbox') else bbox[2]

            if text_width <= max_width:
                current_line = test_line
            else:
                if current_line:
                    lines.append(current_line)
                current_line = char

        if current_line:
            lines.append(current_line)

        # 如果只有一行但太长，强制分割
        if len(lines) == 1 and len(lines[0]) > 15:
            mid = len(lines[0]) // 2
            lines = [lines[0][:mid], lines[0][mid:]]

        return lines

    def _create_prompt(self, text):
        """
        从中文创建英文提示词

        支持角色：小明、小红、老师、小猫、小狗等
        新角色需要在此字典中添加
        """
        prompt_mapping = {
            "小明": "1boy, Chinese student, school uniform, young boy",
            "小红": "1girl, Chinese student, school uniform, young girl",
            "小猫": "small orange cat, cute animal, kitten",
            "小狗": "small dog, cute animal, puppy",
            "老师": "teacher, adult, kind expression, professional",
            "学校": "school building, classroom, educational setting",
            "回家": "walking home, street, road, neighborhood",
            "夕阳": "sunset, golden hour, long shadow, warm light",
            "路上": "street, road, walking path, outdoors",
            "吃饭": "eating, food, meal, dining",
            "睡觉": "sleeping, bed, resting, night",
            "学习": "studying, books, desk, education",
            "玩耍": "playing, happy, fun, recreation",
            "面包": "bread, food, eating",
            "温暖": "warm, happy, emotional, heartwarming",
            "橙色": "orange color",
            "可怜": "sad, pitiful, begging expression",
            "尾巴": "tail, animal tail",
            "微笑": "smiling, happy expression",
            "摸头": "petting, gentle touch, caring",
            "书包": "backpack, school bag",
            "背着": "carrying, wearing backpack"
        }

        # 基础提示词
        base_prompt = "masterpiece, best quality, anime style, detailed, vibrant colors, high resolution"

        # 添加关键词
        additional_keywords = []
        for key, value in prompt_mapping.items():
            if key in text:
                additional_keywords.append(value)

        if additional_keywords:
            return f"{base_prompt}, {', '.join(additional_keywords)}"
        else:
            return f"{base_prompt}, {text}"


def interactive_comic_creator():
    """
    🎨 交互式漫画创建器

    使用流程：
    1. 输入故事，用句号分隔不同场景
    2. 每个句号会自动生成一张图片
    3. 建议每个场景15-30字，描述具体明确
    4. 建议一次生成3-8个场景

    示例输入：
    小明走在回家的路上。他看到一只小猫。他喂小猫吃面包。
    """
    print("🎨 交互式漫画创建器")
    print("=" * 50)

    # 创建生成器实例
    generator = ConfigurableComicGenerator()

    # 加载模型
    print("🔄 正在初始化...")
    if not generator.load_model():
        print("❌ 初始化失败，无法继续")
        return

    # 输入漫画内容
    print("\n📝 请输入您的漫画故事（用句号分隔不同场景）:")
    print("💡 提示：每个句号生成一张图片，建议3-8个场景")
    print("📝 示例: 小明走在回家的路上。他看到一只小猫。他喂小猫吃面包。")
    print("输入 'quit' 结束输入")
    print("-" * 50)

    story_lines = []
    while True:
        try:
            line = input("> ")
            if line.lower() == 'quit':
                break
            if line.strip():
                story_lines.append(line.strip())
        except KeyboardInterrupt:
            print("\n⏹️ 用户中断输入")
            return
        except EOFError:
            break

    if not story_lines:
        print("❌ 没有输入内容")
        return

    full_story = ' '.join(story_lines)

    # 检查场景数量
    scene_count = full_story.count('。')
    if scene_count > 10:
        print(f"⚠️  检测到{scene_count}个场景，建议控制在10个以内")

    # 配置输出
    output_name = input("\n📁 输入输出文件夹名称 (默认: my_comic): ") or "my_comic"

    print(f"\n🚀 开始生成漫画...")
    print(f"故事长度: {len(full_story)} 字符")
    print(f"场景数量: {scene_count} 个")
    print(f"输出到: {output_name}")
    print("-" * 50)

    # 生成漫画
    success = generator.generate_from_text(full_story, output_name)

    if success:
        print(f"\n🎉 漫画生成完成！")
        print(f"📁 文件保存在: {output_name}/ 文件夹中")

        # 尝试自动打开文件夹
        try:
            if os.name == 'nt':  # Windows
                os.startfile(output_name)
            elif os.name == 'posix':  # macOS 或 Linux
                import sys
                if sys.platform == "darwin":
                    os.system(f'open "{output_name}"')
                else:
                    os.system(f'xdg-open "{output_name}"')
            print("✅ 已自动打开输出文件夹")
        except:
            print(f"📁 请手动打开文件夹查看: {output_name}")
    else:
        print("❌ 漫画生成失败")


if __name__ == "__main__":
    interactive_comic_creator()